# -*- coding: utf-8 -*-
"""
预处理 -> 逐年“回归后插值”（默认多项式，AICc选阶）
-> 连续地(2025-03-01 到 2026-07-01) 训练与预测
-> 模型二选一： (A) SARIMAX[差分+MA] + Fourier季节 + 线性时间趋势(外生)，
               (B) UCM(局部线性趋势+季节)
按AIC自动挑选；强制“连续预测”避免 2025/2026 完全一致。
"""

# ============ 参数区（你可改） ============
INPUT_XLSX = "P1sortdata.xlsx"   # Excel 文件名
SHEET_NAME = 0                   # 工作表 index 或 名称

# 等间隔分辨率（天）：你要求默认 4 天，也可改为 2/3/5…；若想自动取全局最小正间隔，写 "auto"
RESOLUTION_DAYS = 3

# 逐年回归插值方式: "poly" 或 "loess"
REGRESSION_KIND = "poly"
POLY_DEGREE_CANDIDATES = (1, 2, 3)  # AICc 自动选
LOESS_FRAC = 0.6                    # loess 平滑强度

# Fourier 阶数
FOURIER_K = 3

# 是否在外生特征中加入线性时间趋势（强烈建议 True）
USE_TIME_TREND = True

# 目标预测区间（含端点）
TARGET_START_MMDD = (3, 1)
TARGET_END_MMDD   = (7, 1)

OUTPUT_DIR = "./output_p1"
# =======================================

import os, warnings, numpy as np, pandas as pd
from datetime import timedelta
warnings.filterwarnings("ignore")
os.makedirs(OUTPUT_DIR, exist_ok=True)

from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import Ridge
from sklearn.pipeline import Pipeline
from scipy.interpolate import interp1d

from statsmodels.tsa.statespace.sarimax import SARIMAX
from statsmodels.tsa.statespace.structural import UnobservedComponents
try:
    from statsmodels.nonparametric.smoothers_lowess import lowess
except Exception:
    lowess = None

# ------------ 工具函数 ------------
def _infer_cols(df: pd.DataFrame):
    date_candidates  = ["日期","date","时间","Date","time","Time"]
    value_candidates = ["就业率","就业率(%)","就业率（%）","y","value","rate"]
    dcol = None
    for c in df.columns:
        try:
            pd.to_datetime(df[c], errors="raise")
            dcol = c; break
        except Exception:
            continue
    if dcol is None:
        for c in date_candidates:
            if c in df.columns: dcol = c; break
    vcol = None
    for c in df.columns:
        if c == dcol: continue
        s = pd.to_numeric(df[c], errors="coerce")
        if s.notna().sum() >= max(5, int(0.3*len(s))):
            vcol = c
            if any(k in str(c) for k in ["就","rate","value","%"]): break
    if vcol is None:
        for c in value_candidates:
            if c in df.columns: vcol = c; break
    if dcol is None or vcol is None:
        raise ValueError(f"无法识别日期/数值列，请检查列名：{df.columns.tolist()}")
    return dcol, vcol

def load_and_clean(path, sheet=0):
    df = pd.read_excel(path, sheet_name=sheet)
    dcol, vcol = _infer_cols(df)
    df = df[[dcol, vcol]].copy()
    df.columns = ["date","y"]
    df["date"] = pd.to_datetime(df["date"], errors="coerce")
    df["y"]    = pd.to_numeric(df["y"], errors="coerce")
    df = df.dropna().drop_duplicates(subset=["date"]).sort_values("date").reset_index(drop=True)
    # 统一到百分比
    if df["y"].median() <= 1.5 and df["y"].max() <= 1.2:
        df["y"] = df["y"] * 100.0
    df["y"] = df["y"].clip(0, 100)
    df["year"] = df["date"].dt.year
    # 全局最小正间隔（天）
    gaps = []
    for _, g in df.groupby("year"):
        d = g["date"].sort_values().diff().dt.days
        pos = d[d > 0]
        if len(pos): gaps.append(int(pos.min()))
    global_min_gap = min(gaps) if gaps else 3
    print(f"[Info] 全局最小正间隔(天)≈ {global_min_gap}")
    return df, global_min_gap

def make_target_grid(year, start_mmdd, end_mmdd, step_days):
    start = pd.Timestamp(year=year, month=start_mmdd[0], day=start_mmdd[1])
    end   = pd.Timestamp(year=year, month=end_mmdd[0],   day=end_mmdd[1])
    return pd.date_range(start, end, freq=f"{int(step_days)}D")

def _aicc(n, sse, k):
    if n <= k + 1: 
        return np.inf
    return n*np.log(max(sse/n, 1e-12)) + 2*k + (2*k*(k+1))/(n - k - 1)

def fit_poly_reg(x, y, deg, alpha=1e-6):
    model = Pipeline([
        ("poly", PolynomialFeatures(degree=deg, include_bias=True)),
        ("ridge", Ridge(alpha=alpha, fit_intercept=False, random_state=0))
    ])
    model.fit(x.reshape(-1,1), y)
    yhat = model.predict(x.reshape(-1,1))
    sse = float(np.sum((y - yhat)**2))
    k = deg + 1
    return model, _aicc(len(y), sse, k)

def per_year_regression_interpolate(df, step_days, kind="poly"):
    """每年先做回归（度数AICc选/或LOESS），再在 3/1~7/1 的等间隔网格取值"""
    out = []
    for y, g in df.groupby("year"):
        start = pd.Timestamp(year=y, month=TARGET_START_MMDD[0], day=TARGET_START_MMDD[1])
        end   = pd.Timestamp(year=y, month=TARGET_END_MMDD[0],   day=TARGET_END_MMDD[1])
        g_win = g[(g["date"] >= start) & (g["date"] <= end)].copy()
        g_use = g_win if len(g_win) >= 2 else g.copy()

        grid = make_target_grid(y, TARGET_START_MMDD, TARGET_END_MMDD, step_days)

        if len(g_use) < 2:
            # 回退：线性/最近邻
            if len(g) >= 2:
                xi = g["date"].astype("int64").values
                yi = g["y"].values
                try:
                    f = interp1d(xi, yi, kind="linear", fill_value="extrapolate")
                except Exception:
                    f = interp1d(xi, yi, kind="nearest", fill_value="extrapolate")
                yg = f(grid.astype("int64").values)
            else:
                yg = np.repeat(g["y"].iloc[0], len(grid))
            out.append(pd.DataFrame({"date": grid, "y": np.clip(yg,0,100), "year": y}))
            continue

        x = g_use["date"].dt.dayofyear.values.astype(float)
        yv = g_use["y"].values.astype(float)
        xg = grid.dayofyear.values.astype(float)

        if kind == "loess" and lowess is not None and len(g_use) >= 5:
            sm = lowess(yv, x, frac=LOESS_FRAC, it=0, return_sorted=True)
            f = interp1d(sm[:,0], sm[:,1], kind="linear", fill_value="extrapolate")
            yg = f(xg)
        else:
            best = (None, np.inf, None)
            for deg in POLY_DEGREE_CANDIDATES:
                try:
                    m, aicc = fit_poly_reg(x, yv, deg)
                    if aicc < best[1]:
                        best = (deg, aicc, m)
                except Exception:
                    continue
            model = best[2] if best[2] is not None else Ridge().fit(x.reshape(-1,1), yv)
            yg = model.predict(xg.reshape(-1,1))

        out.append(pd.DataFrame({"date": grid, "y": np.clip(yg,0,100), "year": y}))
    if not out:
        raise RuntimeError("逐年回归插值后为空。")
    return pd.concat(out, ignore_index=True).sort_values(["year","date"]).reset_index(drop=True)

# ---- 特征构造（含线性时间趋势，保证跨年外推差异） ----
def make_features(dates, K=3, base_date=None, use_trend=True):
    dates = pd.to_datetime(dates)
    idx = pd.DatetimeIndex(dates)
    doy = idx.dayofyear.values.astype(float)
    X = {"intercept": np.ones_like(doy)}
    for k in range(1, K+1):
        X[f"sin_{k}"] = np.sin(2*np.pi*k * doy / 365.25)
        X[f"cos_{k}"] = np.cos(2*np.pi*k * doy / 365.25)
    if use_trend and base_date is not None:
        # 连续时间趋势：以训练起点为零点，单位=年
        trend_years = (idx - pd.Timestamp(base_date)).days.values / 365.25
        X["trend_years"] = trend_years
    return pd.DataFrame(X, index=idx)

def fit_sarimax_with_exog(y, dates, K, base_date, use_trend):
    exog = make_features(dates, K=K, base_date=base_date, use_trend=use_trend)
    model = SARIMAX(
        endog=y,
        exog=exog,
        order=(0,1,1),            # 简洁稳妥
        enforce_stationarity=False,
        enforce_invertibility=False
    )
    res = model.fit(disp=False)
    return res

def fit_ucm(y, period_len):
    model = UnobservedComponents(
        endog=y,
        level="local linear trend",
        seasonal=period_len,
        stochastic_level=True,
        stochastic_seasonal=True
    )
    res = model.fit(disp=False)
    return res

# ---- 训练与“连续预测” ----
def train_and_forecast(interped_df, step_days, fourier_k=3, use_trend=True):
    data = interped_df.sort_values("date").reset_index(drop=True)
    # 训练集：全部历史年（<2025）
    train = data[data["date"].dt.year < 2025].copy()
    if len(train) < 10:
        raise RuntimeError("训练样本过少。")

    y_train = train["y"].values
    d_train = pd.DatetimeIndex(train["date"])
    base_date = d_train.min()

    # 目标整体连续区间：2025-03-01 到 2026-07-01（步长=step_days），一次性预测完
    fut_idx_2025 = make_target_grid(2025, TARGET_START_MMDD, TARGET_END_MMDD, step_days)
    fut_idx_2026 = make_target_grid(2026, TARGET_START_MMDD, TARGET_END_MMDD, step_days)
    full_future_idx = fut_idx_2025.append(fut_idx_2026)

    # 模型A：SARIMAX + Fourier + 线性趋势
    picked = None
    try:
        res_a = fit_sarimax_with_exog(y_train, d_train, K=fourier_k, base_date=base_date, use_trend=use_trend)
        aic_a = res_a.aic
        picked = ("SARIMAX_Fourier+Trend", res_a, aic_a)
        print(f"[Info] SARIMAX_Fourier+Trend AIC={aic_a:.2f}")
    except Exception as e:
        print(f"[Warn] SARIMAX_Fourier+Trend 失败: {e}")

    # 模型B：UCM（局部线性趋势+季节）
    try:
        steps_per_year = len(fut_idx_2025)  # 把一个窗口长度当作“季节周期”
        res_b = fit_ucm(y_train, period_len=steps_per_year)
        aic_b = res_b.aic
        print(f"[Info] UCM(LLT+Seasonal) AIC={aic_b:.2f}")
        if picked is None or aic_b < picked[2]:
            picked = ("UCM", res_b, aic_b)
    except Exception as e:
        print(f"[Warn] UCM 失败: {e}")

    if picked is None:
        raise RuntimeError("模型拟合失败。")

    model_name, res, _ = picked
    print(f"[Info] 采用模型：{model_name}")

    # —— 连续预测 —— #
    if model_name == "SARIMAX_Fourier+Trend":
        X_fut = make_features(full_future_idx, K=fourier_k, base_date=base_date, use_trend=use_trend)
        fc = res.get_forecast(steps=len(full_future_idx), exog=X_fut)
    else:
        fc = res.get_forecast(steps=len(full_future_idx))
    yhat = np.clip(np.asarray(fc.predicted_mean), 0, 100)

    fut_df = pd.DataFrame({"date": full_future_idx, "就业率（%）": yhat})
    fut_df["year"] = fut_df["date"].dt.year

    # 按年拆分输出
    out_2025 = fut_df[fut_df["year"] == 2025].copy()
    out_2026 = fut_df[fut_df["year"] == 2026].copy()
    out_2025.to_csv(os.path.join(OUTPUT_DIR, f"pred_{model_name}_2025_0301_0701_every{int(step_days)}d.csv"),
                    index=False, encoding="utf-8-sig")
    out_2026.to_csv(os.path.join(OUTPUT_DIR, f"pred_{model_name}_2026_0301_0701_every{int(step_days)}d.csv"),
                    index=False, encoding="utf-8-sig")
    fut_df.to_csv(os.path.join(OUTPUT_DIR, f"pred_{model_name}_2025_2026_all_every{int(step_days)}d.csv"),
                  index=False, encoding="utf-8-sig")
    return model_name, out_2025, out_2026

# ------------ 主流程 ------------
def main():
    df, global_min_gap = load_and_clean(INPUT_XLSX, SHEET_NAME)
    df.to_csv(os.path.join(OUTPUT_DIR, "cleaned_all.csv"), index=False, encoding="utf-8-sig")

    # 分辨率
    step_days = max(1, int(global_min_gap if RESOLUTION_DAYS == "auto" else RESOLUTION_DAYS))
    print(f"[Info] 统一分辨率 = 每 {step_days} 天一个点")

    # 逐年回归 -> 栅格取值（仅 3/1~7/1）
    interped = per_year_regression_interpolate(df, step_days, kind=REGRESSION_KIND)
    interped.to_csv(os.path.join(OUTPUT_DIR, f"interpolated_{REGRESSION_KIND}_0301_0701_every{step_days}d.csv"),
                    index=False, encoding="utf-8-sig")
    print("[Info] 已输出逐年“回归后插值”的训练序列。")

    # 训练 + 连续预测（2025整段接着推到2026）
    model_name, pred25, pred26 = train_and_forecast(interped, step_days, fourier_k=FOURIER_K, use_trend=USE_TIME_TREND)
    print("[Done] 2025/2026 预测完成，输出目录：", OUTPUT_DIR)

if __name__ == "__main__":
    main()
